{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to XGBoost Spark with GPU\n",
    "\n",
    "Taxi is an example of XGBoost regressor. This notebook will show you how to load data, train the XGBoost model and use this model to predict \"fare_amount\" of your taxi trip.\n",
    "\n",
    "## Load libraries\n",
    "First load some common libraries will be used by both GPU version and CPU version XGBoost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ml.dmlc.xgboost4j.scala.spark.{XGBoostRegressor, XGBoostRegressionModel}\n",
    "import org.apache.spark.sql.SparkSession\n",
    "import org.apache.spark.ml.evaluation.RegressionEvaluator\n",
    "import org.apache.spark.sql.types.{DoubleType, IntegerType, StructField, StructType}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides CPU version requires some extra libraries, such as:\n",
    "\n",
    "```scala\n",
    "import org.apache.spark.ml.feature.VectorAssembler\n",
    "import org.apache.spark.sql.DataFrame\n",
    "import org.apache.spark.sql.functions._\n",
    "import org.apache.spark.sql.types.FloatType\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the dataset path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dataRoot = /data\n",
       "trainPath = /data/taxi/csv/train/\n",
       "evalPath = /data/taxi/csv/test/\n",
       "transPath = /data/taxi/csv/test/\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "/data/taxi/csv/test/"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// You need to update them to your real paths! The input data files can be the output of taxi-etl jobs, or you can\n",
    "// just use the provided sample datasets upder datasets path. \n",
    "val dataRoot = sys.env.getOrElse(\"DATA_ROOT\", \"/data\")\n",
    "val trainPath = dataRoot + \"/taxi/csv/train/\"\n",
    "val evalPath  = dataRoot + \"/taxi/csv/test/\"\n",
    "val transPath = dataRoot + \"/taxi/csv/test/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the schema of the dataset\n",
    "The Taxi data has 16 columns: 15 features and 1 label. \"fare_amount\" is the label column. The schema will be used to load data in the future. \n",
    "\n",
    "The next block also defines some key parameters used in XGBoost training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "labelName = fare_amount\n",
       "schema = <lazy>\n",
       "featureNames = Array(vendor_id, passenger_count, trip_distance, pickup_longitude, pickup_latitude, rate_code, store_and_fwd, dropoff_longitude, dropoff_latitude, hour, year, month, day, day_of_week, is_weekend)\n",
       "paramMap = <lazy>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<lazy>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val labelName = \"fare_amount\"\n",
    "lazy val schema =\n",
    "  StructType(Array(\n",
    "    StructField(\"vendor_id\", DoubleType),\n",
    "    StructField(\"passenger_count\", DoubleType),\n",
    "    StructField(\"trip_distance\", DoubleType),\n",
    "    StructField(\"pickup_longitude\", DoubleType),\n",
    "    StructField(\"pickup_latitude\", DoubleType),\n",
    "    StructField(\"rate_code\", DoubleType),\n",
    "    StructField(\"store_and_fwd\", DoubleType),\n",
    "    StructField(\"dropoff_longitude\", DoubleType),\n",
    "    StructField(\"dropoff_latitude\", DoubleType),\n",
    "    StructField(labelName, DoubleType),\n",
    "    StructField(\"hour\", DoubleType),\n",
    "    StructField(\"year\", IntegerType),\n",
    "    StructField(\"month\", IntegerType),\n",
    "    StructField(\"day\", DoubleType),\n",
    "    StructField(\"day_of_week\", DoubleType),\n",
    "    StructField(\"is_weekend\", DoubleType)\n",
    "  ))\n",
    "\n",
    "val featureNames = schema.filter(_.name != labelName).map(_.name).toArray\n",
    "\n",
    "lazy val paramMap = Map(\n",
    "  \"num_round\" -> 100\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a new spark session and load data\n",
    "\n",
    "A new spark session should be created to continue all the following spark operations.\n",
    "\n",
    "NOTE: in this notebook, the dependency jars have been loaded when installing toree kernel. Alternatively the jars can be loaded into notebook by [%AddJar magic](https://toree.incubator.apache.org/docs/current/user/faq/). However, there's one restriction for `%AddJar`: the jar uploaded can only be available when `AddJar` is called just after a new spark session is created. Do it as below:\n",
    "\n",
    "```scala\n",
    "import org.apache.spark.sql.SparkSession\n",
    "val spark = SparkSession.builder().appName(\"taxi-GPU\").getOrCreate\n",
    "%AddJar file:/data/libs/rapids-4-spark-XXX.jar\n",
    "%AddJar file:/data/libs/xgboost4j-spark-gpu_2.12-XXX.jar\n",
    "%AddJar file:/data/libs/xgboost4j-gpu_2.12-XXX.jar\n",
    "// ...\n",
    "```\n",
    "\n",
    "##### Please note the new jar \"rapids-4-spark-XXX.jar\" is only needed for GPU version, you can not add it to dependence list for CPU version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "sparkSession = org.apache.spark.sql.SparkSession@6efbc93b\n",
       "reader = org.apache.spark.sql.DataFrameReader@64b8d6da\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "org.apache.spark.sql.DataFrameReader@64b8d6da"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// Build the spark session and data reader as usual\n",
    "val sparkSession = SparkSession.builder().appName(\"taxi-GPU\").getOrCreate\n",
    "val reader = sparkSession.read.option(\"header\", true).schema(schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "trainSet = [vendor_id: double, passenger_count: double ... 14 more fields]\n",
       "evalSet = [vendor_id: double, passenger_count: double ... 14 more fields]\n",
       "transSet = [vendor_id: double, passenger_count: double ... 14 more fields]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[vendor_id: double, passenger_count: double ... 14 more fields]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// Please make sure to change the api to reader.parquet if you load parquet files.\n",
    "val trainSet = reader.csv(trainPath)\n",
    "val evalSet  = reader.csv(evalPath)\n",
    "val transSet = reader.csv(transPath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set XGBoost parameters and build a XGBoostRegressor\n",
    "\n",
    "For CPU version, `num_workers` is recommended being equal to the number of CPU cores, while for GPU version, it should be set to the number of GPUs in Spark cluster.\n",
    "\n",
    "Besides the `tree_method` for CPU version is also different from that for GPU version. Now only \"gpu_hist\" is supported for training on GPU.\n",
    "\n",
    "```scala\n",
    "// difference in parameters\n",
    "  \"num_workers\" -> 12\n",
    "  \"tree_method\" -> \"hist\","
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "xgbParamFinal = Map(num_round -> 100, tree_method -> gpu_hist, num_workers -> 1)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Map(num_round -> 100, tree_method -> gpu_hist, num_workers -> 1)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val xgbParamFinal = paramMap ++ Map(\"tree_method\" -> \"gpu_hist\", \"num_workers\" -> 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "xgbRegressor = xgbr_d36c6f5fd67c\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "xgbr_d36c6f5fd67c"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val xgbRegressor = new XGBoostRegressor(xgbParamFinal)\n",
    "  .setLabelCol(labelName)\n",
    "  .setFeaturesCol(featureNames)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmark and train\n",
    "The object `benchmark` is used to compute the elapsed time of some operations.\n",
    "\n",
    "Training with evaluation dataset is also supported, the same as CPU version's behavior:\n",
    "\n",
    "* Call API `setEvalDataset` after initializing an XGBoostClassifier\n",
    "\n",
    "```scala\n",
    "xgbClassifier.setEvalDataset(evalSet)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "xgbr_d36c6f5fd67c"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xgbRegressor.setEvalDataset(evalSet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "defined object Benchmark\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "object Benchmark {\n",
    "  def time[R](phase: String)(block: => R): (R, Float) = {\n",
    "    val t0 = System.currentTimeMillis\n",
    "    val result = block // call-by-name\n",
    "    val t1 = System.currentTimeMillis\n",
    "    println(\"Elapsed time [\" + phase + \"]: \" + ((t1 - t0).toFloat / 1000) + \"s\")\n",
    "    (result, (t1 - t0).toFloat / 1000)\n",
    "  }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracker started, with env={DMLC_NUM_SERVER=0, DMLC_TRACKER_URI=10.19.183.210, DMLC_TRACKER_PORT=37275, DMLC_NUM_WORKER=1}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "model = xgbr_d36c6f5fd67c\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time [train]: 7.441s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "xgbr_d36c6f5fd67c"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// start training\n",
    "val (model, _) = Benchmark.time(\"train\") {\n",
    "  xgbRegressor.fit(trainSet)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformation and evaluation\n",
    "Here uses `transSet` to evaluate our model and use some key columns to show our predictions. Finally we use `RegressionEvaluator` to calculate an overall `rmse` of our predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time [transform]: 2.134s\n",
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "|    vendor_id|passenger_count|     trip_distance|fare_amount|        prediction|\n",
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "|1.559730423E9|            2.0|0.7000000000000001|        5.0|  5.04693603515625|\n",
      "|1.559730423E9|            3.0|10.700000000000001|       34.0|31.727073669433594|\n",
      "|1.559730423E9|            1.0|               2.3|       10.0| 9.294451713562012|\n",
      "|1.559730423E9|            1.0|               4.4|       16.5| 15.05233097076416|\n",
      "|1.559730423E9|            1.0|               1.5|        7.0| 8.995831489562988|\n",
      "|1.559730423E9|            1.0|               0.8|        7.5| 6.239481449127197|\n",
      "|1.559730423E9|            1.0|               1.2|        5.5| 7.339130401611328|\n",
      "|1.559730423E9|            1.0|               3.0|        2.5|13.403449058532715|\n",
      "| 4.52563162E8|            1.0|2.3399999999999994|        9.5| 9.672189712524414|\n",
      "| 4.52563162E8|            1.0|              3.17|       12.0|11.674100875854492|\n",
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "Elapsed time [evaluation]: 0.17s\n",
      "RMSE == 1.9141528880798715\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "prediction = [vendor_id: double, passenger_count: double ... 15 more fields]\n",
       "evaluator = RegressionEvaluator: uid=regEval_547b9abc7a3b, metricName=rmse, throughOrigin=false\n",
       "rmse = 1.9141528880798715\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "1.9141528880798715"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// start transform\n",
    "val (prediction, _) = Benchmark.time(\"transform\") {\n",
    "  val ret = model.transform(transSet).cache()\n",
    "  ret.foreachPartition((_: Iterator[_]) => ())\n",
    "  ret\n",
    "}\n",
    "prediction.select(\"vendor_id\", \"passenger_count\", \"trip_distance\", labelName, \"prediction\").show(10)\n",
    "val evaluator = new RegressionEvaluator().setLabelCol(labelName)\n",
    "val (rmse, _) = Benchmark.time(\"evaluation\") {\n",
    "  evaluator.evaluate(prediction)\n",
    "}\n",
    "println(s\"RMSE == $rmse\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the model to disk and load model\n",
    "Save the model to disk and then load it to memory. After that use the loaded model to do a new prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time [transform2]: 0.025s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "modelFromDisk = xgbr_d36c6f5fd67c\n",
       "results2 = [vendor_id: double, passenger_count: double ... 15 more fields]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "|    vendor_id|passenger_count|     trip_distance|fare_amount|        prediction|\n",
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "|1.559730423E9|            2.0|0.7000000000000001|        5.0|  5.04693603515625|\n",
      "|1.559730423E9|            3.0|10.700000000000001|       34.0|31.727073669433594|\n",
      "|1.559730423E9|            1.0|               2.3|       10.0| 9.294451713562012|\n",
      "|1.559730423E9|            1.0|               4.4|       16.5| 15.05233097076416|\n",
      "|1.559730423E9|            1.0|               1.5|        7.0| 8.995831489562988|\n",
      "+-------------+---------------+------------------+-----------+------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[vendor_id: double, passenger_count: double ... 15 more fields]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.write.overwrite.save(dataRoot + \"/model/taxi\")\n",
    "\n",
    "val modelFromDisk = XGBoostRegressionModel.load(dataRoot + \"/model/taxi\")\n",
    "val (results2, _) = Benchmark.time(\"transform2\") {\n",
    "  modelFromDisk.transform(transSet)\n",
    "}\n",
    "results2.select(\"vendor_id\", \"passenger_count\", \"trip_distance\", labelName, \"prediction\").show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "sparkSession.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "XGBoost4j-Spark - Scala",
   "language": "scala",
   "name": "XGBoost4j-Spark_scala"
  },
  "language_info": {
   "codemirror_mode": "text/x-scala",
   "file_extension": ".scala",
   "mimetype": "text/x-scala",
   "name": "scala",
   "pygments_lexer": "scala",
   "version": "2.12.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
